from __init__ import *
"""目标检测后处理模块"""
import numpy as np
from typing import Tuple, Union, List
from utils.geometry import GeometryUtils, CoordinateConverter

class DetectorPostProcessor:
    """目标检测后处理器"""
    
    def __init__(self, 
                 conf_threshold: float = 0.25,
                 iou_threshold: float = 0.5,
                 num_classes: int = 80):
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        self.num_classes = num_classes
    
    def process_raw_output(self, 
                          output: Union[np.ndarray, List[np.ndarray]],
                          img_shape: Tuple[int, int],
                          scale_factor: float,
                          pad_w: float,
                          pad_h: float) -> np.ndarray:
        """
        处理原始推理输出
        
        Args:
            output: 推理输出
            img_shape: 原始图像尺寸 (height, width)
            scale_factor: 缩放因子
            pad_w: 水平填充
            pad_h: 垂直填充
            
        Returns:
            检测结果 [x1, y1, x2, y2, score, class_id]
        """
        # 检测输出格式并选择合适的后处理方法
        if self._is_nms_output(output):
            return self._process_nms_output(
                output, img_shape, scale_factor, pad_w, pad_h
            )
        elif self._is_raw_output(output):
            return self._process_raw_output(
                output, img_shape, scale_factor, pad_w, pad_h
            )
        else:
            raise ValueError(f"Unknown output format: {self._get_output_info(output)}")
    
    def _is_nms_output(self, output) -> bool:
        """判断是否为NMS输出格式"""
        return (
            isinstance(output, (list, tuple)) and
            len(output) == 2 and
            hasattr(output[1], "shape") and
            output[1].ndim >= 2 and
            output[1].shape[-1] == 6
        )
    
    def _is_raw_output(self, output) -> bool:
        """判断是否为原始输出格式"""
        return (
            isinstance(output, np.ndarray) and
            output.shape[-2:] == (84, 8400)
        )
    
    def _get_output_info(self, output) -> str:
        """获取输出信息用于调试"""
        if isinstance(output, (list, tuple)):
            return f"List/Tuple with shapes: {[getattr(o, 'shape', None) for o in output]}"
        elif isinstance(output, np.ndarray):
            return f"Array with shape: {output.shape}"
        else:
            return f"Type: {type(output)}"
    
    def _process_nms_output(self, 
                           output: Tuple[np.ndarray, np.ndarray],
                           img_shape: Tuple[int, int],
                           scale_factor: float,
                           pad_w: float,
                           pad_h: float) -> np.ndarray:
        """处理NMS引擎输出"""
        valid_count, detections = output
        
        # 处理维度
        if detections.ndim == 3:
            detections = detections[0]
            valid_count = valid_count[0, 0] if valid_count.ndim == 2 else valid_count[0]
        elif detections.ndim == 2:
            valid_count = valid_count[0] if valid_count.ndim > 0 else valid_count
        else:
            raise ValueError(f"Unknown detections shape: {detections.shape}")
        
        num_detections = int(valid_count)
        if num_detections == 0:
            return np.zeros((0, 6))
        
        detections = detections[:num_detections]
        
        # 坐标反变换
        x1 = (detections[:, 0] - pad_w) / scale_factor
        y1 = (detections[:, 1] - pad_h) / scale_factor
        x2 = (detections[:, 2] - pad_w) / scale_factor
        y2 = (detections[:, 3] - pad_h) / scale_factor
        
        # 裁剪到图像边界
        x1 = np.clip(x1, 0, img_shape[1])
        x2 = np.clip(x2, 0, img_shape[1])
        y1 = np.clip(y1, 0, img_shape[0])
        y2 = np.clip(y2, 0, img_shape[0])
        
        scores = detections[:, 4]
        classes = detections[:, 5]
        
        # 置信度过滤
        mask = scores >= self.conf_threshold
        if np.sum(mask) == 0:
            return np.zeros((0, 6))
        
        return np.stack([x1, y1, x2, y2, scores, classes], axis=1)[mask]
    
    def _process_raw_output(self, 
                           output: np.ndarray,
                           img_shape: Tuple[int, int],
                           scale_factor: float,
                           pad_w: float,
                           pad_h: float) -> np.ndarray:
        """处理原始模型输出"""
        if output.ndim == 3:
            output = output[0]
        
        # 转置为 [num_anchors, num_features]
        predictions = output.transpose(1, 0)  # [8400, 84]
        
        # 分离坐标和分类信息
        boxes = predictions[:, :4]  # [8400, 4]
        class_scores = predictions[:, 4:]  # [8400, 80]
        
        # 获取最大分类得分和对应类别
        class_ids = np.argmax(class_scores, axis=1)
        max_scores = np.max(class_scores, axis=1)
        
        # 置信度过滤
        mask = max_scores >= self.conf_threshold
        boxes = boxes[mask]
        max_scores = max_scores[mask]
        class_ids = class_ids[mask]
        
        if len(boxes) == 0:
            return np.zeros((0, 6))
        
        # 坐标转换 (中心点格式转左上右下格式)
        x_center, y_center, width, height = boxes.T
        x1 = (x_center - width / 2 - pad_w) / scale_factor
        y1 = (y_center - height / 2 - pad_h) / scale_factor
        x2 = (x_center + width / 2 - pad_w) / scale_factor
        y2 = (y_center + height / 2 - pad_h) / scale_factor
        
        # 裁剪到图像边界
        x1 = np.clip(x1, 0, img_shape[1])
        x2 = np.clip(x2, 0, img_shape[1])
        y1 = np.clip(y1, 0, img_shape[0])
        y2 = np.clip(y2, 0, img_shape[0])
        
        # 组合边界框
        boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)
        
        # 应用NMS
        keep_indices = GeometryUtils.nms(boxes_xyxy, max_scores, self.iou_threshold)
        
        # 返回最终结果
        final_boxes = boxes_xyxy[keep_indices]
        final_scores = max_scores[keep_indices]
        final_classes = class_ids[keep_indices]
        
        return np.concatenate([
            final_boxes,
            final_scores[:, None],
            final_classes[:, None]
        ], axis=1)